import cv2
import numpy as np


class VideoReader:
    def __init__(self, video_path):
        """
        初始化 VideoReader 类，打开视频文件。

        参数:
        - video_path: 视频文件的路径
        """
        self.video_path = video_path
        self.cap = cv2.VideoCapture(video_path)
        self.last_frame = None
        if not self.cap.isOpened():
            raise ValueError(f"无法打开视频文件: {video_path}")

    def read_frames(self, num_frames, last_frame_fill=False):
        """
        读取指定数量的帧，并返回一个帧矩阵。

        参数:
        - num_frames: 要读取的帧数

        返回:
        - frames_matrix: 读取到的帧矩阵，形状为 (num_frames, height, width, channels)
        """
        frames = []
        for _ in range(num_frames):
            ret, frame = self.cap.read()
            if ret:
                frames.append(frame)
                self.last_frame = frame
            else:
                break  # 如果没有更多的帧，跳出循环

        # 如果读取的帧数不足，并且需要用最后一帧进行填充
        if last_frame_fill and self.last_frame is not None and len(frames) < num_frames:
            while len(frames) < num_frames:
                frames.append(self.last_frame)

        frames_num = len(frames)
        if frames_num > 0:
            # 将帧列表堆叠成一个 4D 矩阵 (num_frames, height, width, channels)
            frames_matrix = np.stack(frames, axis=0)
        else:
            frames_matrix = None  # 如果没有帧读取到，返回空数组

        return frames_matrix, frames_num

    def get_fps(self):
        fps = self.cap.get(cv2.CAP_PROP_FPS)
        return fps

    def get_resolution(self):
        """
        获取视频的分辨率（宽度和高度）。

        返回:
        - width: 视频的宽度
        - height: 视频的高度
        """
        height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        return height, width

    def release(self):
        """
        释放视频文件的资源。
        """
        self.cap.release()


# 使用示例
if __name__ == "__main__":
    video_path = "data/ori_video/飞书20241014-172537.mp4"
    reader = VideoReader(video_path)

    # 读取5帧
    frames = reader.read_frames(5)

    # 检查是否读取到5帧
    print(f"成功读取了 {len(frames)} 帧")

    # 释放资源
    reader.release()